import org.antlr.v4.runtime.tree.TerminalNode;
import 抽象语法树类.*;
import 抽象语法树类.字面量类.*;
import 抽象语法树类.表达式类.*;
import 抽象语法树类.表达式类.二元表达式类.*;
import 抽象语法树类.语句类.*;

import java.util.ArrayList;

public class 语法树访问器 extends PieBaseVisitor<抽象语法树> {
    @Override
    public 抽象语法树 visitProgram(PieParser.ProgramContext ctx) {
        var 结果 = new 程序节点();
        结果.子节点 = new ArrayList<>();
        for (var node : ctx.children) {
            if (node instanceof TerminalNode) continue; // 跳过EOF
            var 节点 = visit(node);
            if (节点 == null) continue;
            结果.子节点.add(节点);
        }

        return 结果.记录(ctx);
    }

    @Override
    public 抽象语法树 visitStructDefinition(PieParser.StructDefinitionContext ctx) {
        var 结果 = new 结构体定义();

        结果.结构体名 = ctx.name.getText();
        结果.成员 = new ArrayList<>();

        for (var varDef : ctx.vardef()) {
            var 成员 = visit(varDef);
            结果.成员.add((变量定义) 成员);
        }

        return 结果.记录(ctx);
    }


    @Override
    public 抽象语法树 visitFunctionDefinition(PieParser.FunctionDefinitionContext ctx) {
        var 结果 = new 函数定义();

        结果.函数名 = ctx.ID().getText();

        结果.参数 = new ArrayList<>();
        for (var varDef : ctx.vardef()) {
            var 参数 = (变量定义) visit(varDef);
            结果.参数.add(参数);
        }

        结果.代码块 = (代码块) visit(ctx.slist());

        return 结果.记录(ctx);
    }


    @Override
    public 抽象语法树 visitSlist(PieParser.SlistContext ctx) {
        var 结果 = new 代码块();
        结果.语句列表 = new ArrayList<>();

        for (var statement : ctx.statement()) {
            var 语句 = visit(statement);
            if (语句 == null) continue;
            结果.语句列表.add(语句);
        }

        return 结果.记录(ctx);
    }


    @Override
    public 抽象语法树 visitAssignStatement(PieParser.AssignStatementContext ctx) {
        var 结果 = new 赋值语句();
        结果.左值 = visit(ctx.qid());
        结果.右值 = visit(ctx.expr());
        return 结果.记录(ctx);
    }

    @Override
    public 抽象语法树 visitReturnStatement(PieParser.ReturnStatementContext ctx) {
        var 结果 = new 返回语句();
        结果.返回值 = visit(ctx.expr());
        return 结果.记录(ctx);
    }

    @Override
    public 抽象语法树 visitPrintStatement(PieParser.PrintStatementContext ctx) {
        var 结果 = new 打印语句();
        结果.表达式 = visit(ctx.expr());
        return 结果.记录(ctx);
    }

    @Override
    public 抽象语法树 visitIfStatement(PieParser.IfStatementContext ctx) {
        var 结果 = new If语句();
        结果.条件 = visit(ctx.expr());
        结果.真块 = (代码块) visit(ctx.c);

        if (ctx.el != null) {
            结果.假块 = (代码块) visit(ctx.el);
        }

        return 结果.记录(ctx);
    }


    @Override
    public 抽象语法树 visitWhileStatement(PieParser.WhileStatementContext ctx) {
        var 结果 = new While语句();
        结果.条件 = visit(ctx.expr());
        结果.代码块 = (代码块) visit(ctx.slist());
        return 结果.记录(ctx);
    }

    @Override
    public 抽象语法树 visitCallStatement(PieParser.CallStatementContext ctx) {
        return visit(ctx.call());
    }

    @Override
    public 抽象语法树 visitCall(PieParser.CallContext ctx) {
        var 结果 = new 函数调用语句();
        结果.函数名 = ctx.ID().getText();
        结果.参数 = new ArrayList<>();
        for (var expr : ctx.expr()) {
            结果.参数.add(visit(expr));
        }
        return 结果.记录(ctx);
    }


    @Override
    public 抽象语法树 visitMulExpr(PieParser.MulExprContext ctx) {
        var 结果 = new 乘法表达式();
        结果.左值 = visit(ctx.expr(0));
        结果.右值 = visit(ctx.expr(1));
        return 结果.记录(ctx);
    }

    @Override
    public 抽象语法树 visitAddExpr(PieParser.AddExprContext ctx) {
        var 结果 = new 加法表达式();
        结果.左值 = visit(ctx.expr(0));
        结果.右值 = visit(ctx.expr(1));
        return 结果.记录(ctx);
    }

    @Override
    public 抽象语法树 visitSubExpr(PieParser.SubExprContext ctx) {
        var 结果 = new 减法表达式();
        结果.左值 = visit(ctx.expr(0));
        结果.右值 = visit(ctx.expr(1));
        return 结果.记录(ctx);
    }

    @Override
    public 抽象语法树 visitEqualExpr(PieParser.EqualExprContext ctx) {
        var 结果 = new 判等表达式();
        结果.左值 = visit(ctx.expr(0));
        结果.右值 = visit(ctx.expr(1));
        return 结果.记录(ctx);
    }

    @Override
    public 抽象语法树 visitLessThanExpr(PieParser.LessThanExprContext ctx) {
        var 结果 = new 小于表达式();
        结果.左值 = visit(ctx.expr(0));
        结果.右值 = visit(ctx.expr(1));
        return 结果.记录(ctx);
    }


    @Override
    public 抽象语法树 visitAtom(PieParser.AtomContext ctx) {
        if (ctx.INT() != null) {
            var 结果 = new 整数型字面量();
            结果.值 = ctx.INT().getSymbol();
            return 结果.记录(ctx);
        } else if (ctx.CHAR() != null) {
            var 结果 = new 字符型字面量();
            结果.值 = ctx.CHAR().getSymbol();
            return 结果.记录(ctx);
        } else if (ctx.FLOAT() != null) {
            var 结果 = new 浮点型字面量();
            结果.值 = ctx.FLOAT().getSymbol();
            return 结果.记录(ctx);
        } else if (ctx.STRING() != null) {
            var 结果 = new 文本型字面量();
            结果.值 = ctx.STRING().getSymbol();
            return 结果.记录(ctx);
        } else if (ctx.expr() != null) {
            return visit(ctx.expr());
        } else {
            return visit(ctx.getChild(0));
        }
    }


    @Override
    public 抽象语法树 visitInstance(PieParser.InstanceContext ctx) {
        var 结果 = new 实例化表达式();
        结果.类型名 = ctx.sname.getText();
        return 结果.记录(ctx);
    }


    @Override
    public 抽象语法树 visitMemberAccess(PieParser.MemberAccessContext ctx) {
        var 结果 = new 成员访问表达式();
        结果.对象 = visit(ctx.qid(0));
        结果.成员 = visit(ctx.qid(1));
        return 结果.记录(ctx);
    }

    @Override
    public 抽象语法树 visitVarReference(PieParser.VarReferenceContext ctx) {
        var 结果 = new 符号引用();
        结果.符号 = ctx.ID().getText();
        return 结果.记录(ctx);
    }

    @Override
    public 抽象语法树 visitVardef(PieParser.VardefContext ctx) {
        var 结果 = new 变量定义();
        结果.变量名 = ctx.ID().getText();
        return 结果.记录(ctx);
    }
}

